'''
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with self work for additional information
 * regarding copyright ownership.  The ASF licenses self file
 * to you under the Apache License, Version 2.0 (the
 * "License") you may not use self file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http:#www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
'''
from base.graph.FastArrayBunch import FastArrayBunch
from base.graph.TripleStore import TripleStore
from base.graph.impl.MatchPattern import MatchPattern
from base.graph.impl.PatternClassifier import PatternClassifier
from base.java.FastHashSet import FastHashSet
from base.java.HashMap import HashMap
from base.java.exception.IllegalStateException import IllegalStateException
from base.util.iterator.impl.IteratorImpl import IteratorImpl
from base.util.iterator.impl.NiceIterator import NiceIterator

#import org.apache.jena.graph.Triple;
#import org.apache.jena.mem2.collection.FastHashMap;
#import org.apache.jena.mem2.iterator.IteratorOfJenaSets;
#import org.apache.jena.mem2.pattern.PatternClassifier;
#import org.apache.jena.mem2.store.TripleStore;
#import org.apache.jena.util.iterator.ExtendedIterator;
#import org.apache.jena.util.iterator.NiceIterator;
#import org.apache.jena.util.iterator.SingletonIterator;

#import java.util.stream.Stream;
#import java.util.stream.StreamSupport;

'''
 * A triple store that uses hash tables to map from nodes to triple bunches.
 * <p>
 * Inner structure:
 * - three:@link FastHashMap, one for each node (subject, predicate, object) in the triple
 * - each map maps from a node to a:@link FastTripleBunch
 * - for up to 16 triples with the same subject, the bunch is an:@link FastArrayBunch, otherwise it is
 * a:@link FastHashedTripleBunch
 * - for up to 32 triples with the same predicate and object, the bunch is an:@link FastArrayBunch, otherwise it is
 * a:@link FastHashedTripleBunch
 * <p>
 * Other optimizations:
 * - each triple is added to three:@link FastTripleBunches. To avoid the overhead of calculating the hash code three
 * times, the hash code is calculated once and passed to the:@link FastTripleBunches.
 * - the different sizes for the:@link FastArrayBunches are chosen:
 * - to avoid the memory-overhead of:@link FastHashedTripleBunchfor a small number of triples
 * - the subject bunch is smaller than the predicate and object bunches, because the subject is typically used for
 * #contains operations, which are faster for:@link FastHashedTripleBunches.
 * - the predicate and object bunches are the same size, because they are typically used for #find operations, which
 * typically do not answer #contains operations. Making them much larger might slow down #remove operations on the
 * graph.
 * - "ANY_PRE_OBJ" matches primarily use the object bunch unless the size of the bunch is larger than 400 triples. In
 * that case, there is a secondary lookup in the predicate bunch. Then the smaller bunch is used for the lookup.
 * Especially for RDF graphs there are some very common object nodes like "True"/"False" or "0"/"1". In that case,
 * a secondary lookup in the predicate bunch might be faster than a primary lookup in the object bunch.
 * - FastTripleStore#contains uses:@link org.apache.jena.mem2.store.fast.FastTripleBunch#anyMatchRandomOrderfor
 * "ANY_PRE_OBJ" lookups. This is only faster than:@link org.apache.jena.mem2.collection.JenaMapSetCommon#anyMatch
 * if there are many matches and the set is ordered in an unfavorable way. "ANY_PRE_OBJ" matches usually fall into
 * self category. This optimization was only needed because the:@link FastTripleBunches does not use the random
 * order of the triples in #anyMatch but the ordered dense array of triples, which is faster if there are only a few
 * matches.
 * - for the FastArrayBunches, the equals method of the triple is not called. Instead, only the two nodes that are
 * not part of the key of the containing map are compared.
'''
# protected static final int
THRESHOLD_FOR_SECONDARY_LOOKUP = 400
# protected static final int
MAX_ARRAY_BUNCH_SIZE_SUBJECT = 16
# protected static final int
MAX_ARRAY_BUNCH_SIZE_PREDICATE_OBJECT = 32


class ArrayBunchWithSameSubject(FastArrayBunch):

    # @Override
    # param final Triple a, final Triple b
    # result boolean
    def areEqual(self, a, b):
        return a.getPredicate().equals(b.getPredicate()) and a.getObject().equals(b.getObject())


class ArrayBunchWithSamePredicate(FastArrayBunch):
    # @Override
    # param final Triple a, final Triple b
    # result boolean
    def areEqual(self, a, b):
        return a.getSubject().equals(b.getSubject()) and a.getObject().equals(b.getObject())


class ArrayBunchWithSameObject(FastArrayBunch):
    # @Override
    # param final Triple a, final Triple b
    # result boolean
    def areEqual(self, a, b):
        return a.getSubject().equals(b.getSubject()) and a.getPredicate().equals(b.getPredicate())


class FastTripleStore(TripleStore):

    # final FastHashedBunchMap
    subjects = HashMap()
    # final FastHashedBunchMap
    predicates = HashMap()
    # final FastHashedBunchMap
    objects = HashMap()

    def __init__(self):
        self.size = 0

    #@Override
    # param Triple triple
    # result void
    def add(self, triple):
        # final int
        hashCodeOfTriple = triple.hashCode()
        # final boolean 
        added = False
        sBunch = self.subjects.get(triple.getSubject())
        if (sBunch is None):
            #sBunch = ArrayBunchWithSameSubject()
            sBunch = ArrayBunchWithSameSubject()
            sBunch.addUnchecked(triple)
            self.subjects.put(triple.getSubject(), sBunch)
            added = True
        else:
            if (sBunch.isArray() and sBunch.size == MAX_ARRAY_BUNCH_SIZE_SUBJECT):
                #sBunch = FastHashedTripleBunch(sBunch)
                sBunch = FastHashSet(sBunch)
                self.subjects.put(triple.getSubject(), sBunch)
            
            added = sBunch.tryAdd(triple, hashCodeOfTriple)
        
        if (added):
            self.size += 1
            # var
            pBunch = self.predicates.computeIfAbsent(triple.getPredicate(), ArrayBunchWithSamePredicate())
            # print('type pBunch', type(pBunch))
            if (pBunch.isArray() and pBunch.size == MAX_ARRAY_BUNCH_SIZE_PREDICATE_OBJECT):
                pBunch = FastHashSet(pBunch)
                self.predicates.put(triple.getPredicate(), pBunch)

            pBunch.addUnchecked(triple)
            # var
            oBunch = self.objects.computeIfAbsent(triple.getObject(), ArrayBunchWithSameObject())
            if oBunch.isArray() and oBunch.size == MAX_ARRAY_BUNCH_SIZE_PREDICATE_OBJECT:
                oBunch = FastHashSet(oBunch)
                self.objects.put(triple.getObject(), oBunch)
            
            oBunch.addUnchecked(triple)
        
    

    #@Override
    # param Triple triple
    # result void
    def remove(self, triple):
        # final int
        hashCodeOfTriple = triple.hashCode()
        # final var 
        sBunch = self.subjects.get(triple.getSubject())
        if (sBunch is None):
            return

        if (sBunch.tryRemove(triple, hashCodeOfTriple)):
            if (sBunch.isEmpty()):
                self.subjects.removeUnchecked(triple.getSubject())
            
            # final var
            pBunch = self.predicates.get(triple.getPredicate())
            pBunch.removeUnchecked(triple, hashCodeOfTriple)
            if (pBunch.isEmpty()):
                self.predicates.removeUnchecked(triple.getPredicate())
            
            # final var
            oBunch = self.objects.get(triple.getObject())
            oBunch.removeUnchecked(triple, hashCodeOfTriple)
            if (oBunch.isEmpty()):
                self.objects.removeUnchecked(triple.getObject())
            
            self.size -= 1
        
    

    #@Override
    # param 
    # result void
    def clear(self):
        self.subjects.clear()
        self.predicates.clear()
        self.objects.clear()
        size = 0
    

    #@Override
    # param 
    # result int
    def countTriples(self):
        return self.size
    

    #@Override
    # param 
    # result boolean
    def isEmpty(self):
        return self.size == 0
    

    #@Override
    # param Triple tripleMatch
    # result boolean
    def contains(self, tripleMatch):
        classifyResult = PatternClassifier.classify(tripleMatch)
        #switch (PatternClassifier.classify(tripleMatch)):

        if classifyResult == MatchPattern.SUB_PRE_OBJ:
            # final var
            triples = self.subjects.get(tripleMatch.getSubject())
            if (triples is None):
                return False

            return triples.containsKey(tripleMatch)


        elif classifyResult == MatchPattern.SUB_PRE_ANY:
            # final var
            triplesBySubject = self.subjects.get(tripleMatch.getSubject())
            if (triplesBySubject is None):
                return False

            return triplesBySubject.anyMatch(lambda t:tripleMatch.getPredicate().equals(t.getPredicate()))


        elif classifyResult == MatchPattern.SUB_ANY_OBJ:
            # final var
            triplesBySubject = self.subjects.get(tripleMatch.getSubject())
            if (triplesBySubject is None):
                return False

            return triplesBySubject.anyMatch(lambda t:tripleMatch.getObject().equals(t.getObject()))


        elif classifyResult == MatchPattern.SUB_ANY_ANY:
            return self.subjects.containsKey(tripleMatch.getSubject())

        elif classifyResult == MatchPattern.ANY_PRE_OBJ:
            # final var
            triplesByObject = self.objects.get(tripleMatch.getObject())
            if (triplesByObject is None):
                return False

            # Optimization for typical RDF data, where there may be common values like "0" or "False"/"True".
            # In self case, there may be many matches but due to the ordered nature of FastHashBase,
            # the same predicates are often grouped together. If they are at the beginning of the bunch,
            # we can avoid the linear scan of the bunch. This is a common elif classifyResult == for RDF data.
            # #anyMatchRandomOrder is a bit slower if the predicate is not found than #anyMatch, but not by much.
            if (triplesByObject.size > THRESHOLD_FOR_SECONDARY_LOOKUP):
                # final var
                triplesByPredicate = self.predicates.get(tripleMatch.getPredicate())
                if (triplesByPredicate is None):
                    return False

                if (triplesByPredicate.size < triplesByObject.size):
                    return triplesByPredicate.anyMatchRandomOrder(lambda t:tripleMatch.getObject().equals(t.getObject()))


            return triplesByObject.anyMatchRandomOrder(lambda t:tripleMatch.getPredicate().equals(t.getPredicate()))


        elif classifyResult == MatchPattern.ANY_PRE_ANY:
            return self.predicates.containsKey(tripleMatch.getPredicate())

        elif classifyResult == MatchPattern.ANY_ANY_OBJ:
            return self.objects.containsKey(tripleMatch.getObject())

        elif classifyResult == MatchPattern.ANY_ANY_ANY:
            return not self.isEmpty()

        else:
            raise IllegalStateException("Unexpected value: " + PatternClassifier.classify(tripleMatch))
        

    #@Override
    # param 
    # result Stream<Triple>
    # def stream(self):
    #     return StreamSupport.stream(subjects.valueSpliterator(), False)
    #             .flatMap(bunch -> StreamSupport.stream(bunch.keySpliterator(), False))
    

    #@Override
    # param Triple tripleMatch
    # result Stream<Triple>
    # def stream(self, tripleMatch):
    #     classifyResult = PatternClassifier.classify(tripleMatch)
    #     # switch (PatternClassifier.classify(tripleMatch)):
    #
    #     if classifyResult == MatchPattern.SUB_PRE_OBJ:
    #         # final var
    #         triples = self.subjects.get(tripleMatch.getSubject())
    #         if (triples is None):
    #             return Stream.empty()
    #
    #         return triples.containsKeyStream.of(tripleMatch) if (tripleMatch) else Stream.empty()
    #
    #
    #     elif classifyResult == MatchPattern.SUB_PRE_ANY:
    #         # final var
    #         triplesBySubject = self.subjects.get(tripleMatch.getSubject())
    #         if (triplesBySubject is None):
    #             return Stream.empty()
    #
    #         return triplesBySubject.keyStream().filter(lambda t:tripleMatch.getPredicate().equals(t.getPredicate()))
    #
    #
    #     elif classifyResult == SUB_ANY_OBJ:
    #         # final var
    #         triplesBySubject = self.subjects.get(tripleMatch.getSubject())
    #         if (triplesBySubject is None):
    #             return Stream.empty()
    #
    #         return triplesBySubject.keyStream().filter(lambda t:tripleMatch.getObject().equals(t.getObject()))
    #
    #
    #     elif classifyResult == SUB_ANY_ANY:
    #         # final var
    #         triples = self.subjects.get(tripleMatch.getSubject())
    #         return Stream.empty() if triples is None else triples.keyStream()
    #
    #
    #     elif classifyResult == ANY_PRE_OBJ:
    #         # final var
    #         triplesByObject = self.objects.get(tripleMatch.getObject())
    #         if (triplesByObject is None):
    #             return Stream.empty()
    #
    #         if (triplesByObject.size > THRESHOLD_FOR_SECONDARY_LOOKUP):
    #             # final var
    #             triplesByPredicate = self.predicates.get(tripleMatch.getPredicate())
    #             if (triplesByPredicate is None):
    #                 return Stream.empty()
    #
    #             if (triplesByPredicate.size < triplesByObject.size):
    #                 return triplesByPredicate.keyStream().filter(lambda t:tripleMatch.getObject().equals(t.getObject()))
    #
    #
    #         return triplesByObject.keyStream().filter(lambda t:tripleMatch.getPredicate().equals(t.getPredicate()))
    #
    #
    #     elif classifyResult == ANY_PRE_ANY:
    #         # final var
    #         triples = self.predicates.get(tripleMatch.getPredicate())
    #         return Stream.empty() if triples is None else triples.keyStream()
    #
    #
    #     elif classifyResult == ANY_ANY_OBJ:
    #         # final var
    #         triples = self.objects.get(tripleMatch.getObject())
    #         return Stream.empty() if triples is None else triples.keyStream()
    #
    #
    #     elif classifyResult == ANY_ANY_ANY:
    #         return stream()
    #
    #     else:
    #         raise IllegalStateException("Unexpected value: " + PatternClassifier.classify(tripleMatch))
        
    

    #@Override
    # param Triple tripleMatch
    # result ExtendedIterator<Triple>
    def find(self, tripleMatch):
        classifyResult = PatternClassifier.classify(tripleMatch)
        #switch (PatternClassifier.classify(tripleMatch)):

        if classifyResult == MatchPattern.SUB_PRE_OBJ:
            # final var
            triples = self.subjects.get(tripleMatch.getSubject())
            if (triples is None):
                return NiceIterator.emptyIterator()

            return IteratorImpl([tripleMatch]) if triples.containsKey(tripleMatch) else NiceIterator.emptyIterator()


        elif classifyResult == MatchPattern.SUB_PRE_ANY:
            # final var
            triplesBySubject = self.subjects.get(tripleMatch.getSubject())
            if (triplesBySubject is None):
                return NiceIterator.emptyIterator()

            return triplesBySubject.keyIterator().filterKeep(lambda t:tripleMatch.getPredicate().equals(t.getPredicate()))


        elif classifyResult == MatchPattern.SUB_ANY_OBJ:
            # final var
            triplesBySubject = self.subjects.get(tripleMatch.getSubject())
            if (triplesBySubject is None):
                return NiceIterator.emptyIterator()

            return triplesBySubject.keyIterator().filterKeep(lambda t:tripleMatch.getObject().equals(t.getObject()))


        elif classifyResult == MatchPattern.SUB_ANY_ANY:
            # final var
            triples = self.subjects.get(tripleMatch.getSubject())
            return NiceIterator.emptyIterator() if triples is None else triples.keyIterator()


        elif classifyResult == MatchPattern.ANY_PRE_OBJ:
            # final var
            triplesByObject = self.objects.get(tripleMatch.getObject())
            if (triplesByObject is None):
                return NiceIterator.emptyIterator()

            if (triplesByObject.size > THRESHOLD_FOR_SECONDARY_LOOKUP):
                # final var
                triplesByPredicate = self.predicates.get(tripleMatch.getPredicate())
                if (triplesByPredicate is None):
                    return NiceIterator.emptyIterator()

                if (triplesByPredicate.size < triplesByObject.size):
                    return triplesByPredicate.keyIterator().filterKeep(lambda t:tripleMatch.getObject().equals(t.getObject()))


            return triplesByObject.keyIterator().filterKeep(lambda t:tripleMatch.getPredicate().equals(t.getPredicate()))


        elif classifyResult == MatchPattern.ANY_PRE_ANY:
            # final var
            triples = self.predicates.get(tripleMatch.getPredicate())
            return NiceIterator.emptyIterator() if triples is None else triples.keyIterator()


        elif classifyResult == MatchPattern.ANY_ANY_OBJ:
            # final var
            triples = self.objects.get(tripleMatch.getObject())
            return NiceIterator.emptyIterator() if triples is None else triples.keyIterator()


        elif classifyResult == MatchPattern.ANY_ANY_ANY:
            #return IteratorOfJenaSets<>(subjects.valueIterator())
            return IteratorImpl(self.subjects)
        else:
            raise IllegalStateException("Unexpected value: " + PatternClassifier.classify(tripleMatch))
        
    


        
    

